#!/opt/local/bin/python

import matplotlib.pyplot as plt
from scipy import sin, cos, tan, arctan, sqrt, pi, exp
from numpy import array, abs, max, min, isnan, nan, arange, logspace
from plotcsv import plotcsv
from globals import c_light as c

class cavity:
    def __init__(self, f=704.42e6, delF=nan, I_b0=50e-3, I_b=0, V_cav=18e6, beamph=0, RoverQ=477, Ql=nan, beta_0=0.9, beta=0.9, Eampfile=nan):
        if Ql<0: raise ValueError('Ql must be a positive number or nan')
        self.f        = f
        self.delF     = delF
        self.beamph   = beamph
        self.RoverQ   = RoverQ
        try:
            self.RQscaling = scale_roverq(beta_0, file=Eampfile)(beta, freq=f)[0]
        except IndexError:
            self.RQscaling = scale_roverq(beta_0, file=Eampfile)(beta, freq=f)
        self.Ql       = Ql
        self.I_b0     = I_b0 * (-1+0j)   # Beam current in negative real direction
        self.I_b      = I_b * (-1+0j)
        self.V_cav    = V_cav * (1+0j) * exp(1j * (beamph*pi/180))
        self.beta_0   = beta_0
        self.beta     = beta
        self.Eampfile = Eampfile

        self.corrVT()
        self.buildVectors()

    def corrVT(self):
        self.V_cav = self.V_cav * sqrt(self.RQscaling)

    def buildVectors(self):
        self.calcQlOpt()
        self.calcFracFreq()
        self.calcTuningAngle()

        self.calcVb()
        self.calcVg()
        self.calcVfor()
        self.calcGenPhase()
        self.calcPbeam()

    def calcQlOpt(self):
        self.Ql0 = abs(self.V_cav) / (self.RoverQ * self.RQscaling * abs(self.I_b0) * cos(self.beamph*pi/180))
        if isnan(self.Ql):
            self.Ql = self.Ql0

    def calcFracFreq(self):
        self.fracFreq0 = -(self.RoverQ * self.RQscaling * abs(self.I_b0) * sin(self.beamph*pi/180)) / (2*abs(self.V_cav))
        if isnan(self.delF):
            self.delF = self.fracFreq0*self.f
            self.fracFreq = self.fracFreq0
        else:
            self.fracFreq  = self.delF/self.f

    def calcTuningAngle(self):
        self.TAng0 = arctan(2 * self.Ql0 * self.fracFreq0)
        self.TAng  = arctan(2 * self.Ql * self.fracFreq)

    def calcVb(self):
        self.V_b = cos(self.TAng) * exp(1j*self.TAng) * self.RoverQ * self.RQscaling * self.Ql * self.I_b
        self.V_b0 = cos(self.TAng) * exp(1j*self.TAng) * self.RoverQ * self.RQscaling * self.Ql * self.I_b0

    def calcVg(self):
        self.V_g  = self.V_cav - self.V_b
        self.V_g0 = self.V_cav - self.V_b0
        prefact = 0.25 * abs(self.V_cav)**2 / (self.RoverQ*self.RQscaling*self.Ql)
        term1 = 1 + (self.RoverQ*self.RQscaling * self.Ql * abs(self.I_b) * cos(self.beamph*(pi/180)) / abs(self.V_cav))
        term2 = 2 * self.fracFreq * self.Ql + (self.RoverQ*self.RQscaling * self.Ql * abs(self.I_b) * sin(self.beamph*(pi/180)) / abs(self.V_cav))
        #print "prefact = %e :: term1 = %e :: term2 = %e" % (prefact, term1, term2)
        self.Pg = prefact * (term1**2 + term2**2)

    def calcGenPhase(self):
        self.Phig = (180/pi) * arctan(self.V_g.imag / self.V_g.real)

    def calcVfor(self):
        if abs(self.TAng)<1e-4:
            amp = abs(self.V_g)/2
        else:
            amp = abs(self.V_g) * (sin(self.TAng)/sin(pi-(2*self.TAng)))
        self.Vfor = amp * (self.V_g / abs(self.V_g)) * exp(1j * -self.TAng)
        self.Vref = self.V_cav-self.Vfor

    def calcPbeam(self):
        self.Pbeam = abs(self.V_cav) * abs(self.I_b) * cos(self.beamph*(pi/180))
        self.Pref = self.Pg - self.Pbeam

    def drawvecs(self, fig):
        ax = fig.add_subplot(111, aspect='equal')
        if not self.I_b==0:
            ax.plot(([0, self.I_b.real*1e7]), ([0, self.I_b.imag*1e7]), '-r', label='I_b (arb. units)')
            ax.plot(([0, self.V_b.real]), ([0, self.V_b.imag]), '-r', label='V_b')
        ax.plot(([0, self.V_g.real]), ([0, self.V_g.imag]), '-g', label='V_g')
        ax.plot(([0, self.V_cav.real]), ([0, self.V_cav.imag]), '-b', label='V_cav')
        ax.plot(([0, self.Vfor.real]), ([0, self.Vfor.imag]), '-y', label='V_for')
        if abs(self.Vref)/abs(self.Vfor) > 1e-6:
            ax.plot(([self.Vfor.real, self.V_cav.real]), ([self.Vfor.imag, self.V_cav.imag]), '-k', label='V_ref (%0.1f%%)' % (100*abs(self.Vref)/abs(self.Vfor)))

        xscale = ax.get_xlim()
        yscale = ax.get_ylim()
        ax.set_xlim( [1.05*min([xscale[0],yscale[0]]), 1.05*max([xscale[1],yscale[1]])] )
        ax.set_ylim( [1.05*min([xscale[0],yscale[0]]), 1.05*max([xscale[1],yscale[1]])] )

        widtharg = max([xscale[1],yscale[1]]) / 300
        if not self.I_b==0:
            ax.arrow(0, 0, self.I_b.real*1e7, self.I_b.imag*1e7, width=widtharg, color='r')
            ax.arrow(0, 0, self.V_b.real, self.V_b.imag, width=widtharg, color='r')
            ax.arrow(self.V_g.real, self.V_g.imag, self.V_b.real, self.V_b.imag, width=widtharg, color='r')
        ax.arrow(0, 0, self.V_g.real, self.V_g.imag, width=widtharg, color='g')
        ax.arrow(0, 0, self.V_cav.real, self.V_cav.imag, width=widtharg*2, color='b')
        ax.arrow(0, 0, self.Vfor.real, self.Vfor.imag, width=widtharg, color='y')
        if abs(self.Vref)/abs(self.Vfor) > 1e-6:
            ax.arrow(self.Vfor.real, self.Vfor.imag, self.Vref.real*0.95, self.Vref.imag*0.95, width=widtharg, color='k')

        ax.set_xlabel('Real voltage / V')
        ax.set_ylabel('Imaginary voltage / V')
        ax.grid()
        #ax.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
        ax.legend(loc=0)

        text = "Cavity frequency = %0.3f MHz\nR/Q = %0.0f Ohms\nOptimum current = %0.3f A\nActual current = %0.3f A\n"
        text += "Beam phase = %0.3f degrees\nCavity voltage = %0.3f MV\nOptimum Ql = %0.3e\nActual Ql = %0.3e\n"
        text += "Detuning freq = %0.1f Hz\nOptimum = %0.1f Hz\nDetuning angle = %0.3f degrees\nOptimum = %0.3f degrees\n"
        text += "Reflection (voltage) = %0.3f %%\nReflection (power) = %0.3f %%"
        text = text % ((self.f/1e6), self.RoverQ, abs(self.I_b0), abs(self.I_b), self.beamph, abs(self.V_cav)/1e6,
                self.Ql0, self.Ql, (self.fracFreq * self.f), (self.fracFreq0* self.f), (self.TAng*180/pi),
                (self.TAng0*180/pi), (100*abs(self.Vref)/abs(self.Vfor)), (100*(abs(self.Vref)/abs(self.Vfor))**2))

        plt.figtext(0.23,0.5, text, size='x-small')
        self.fig = fig
        self.ax = ax

def scale_roverq(beta_0, file=nan):
    try:
        fieldamp = plotcsv(file, firstrow=3, delimiter=' ')
        def scaling(beta, freq=352.21e6):
            integ = lambda b, a, freq: sum(cos(2*pi*freq*a.x / (1000*b*c)) * a.y)
            norm = integ(beta_0, fieldamp, freq)
            try:
                raw  = integ(beta, fieldamp, freq)
            except ValueError:
                raw = []
                [raw.append(integ(i, fieldamp, freq)) for i in beta]
                raw = array(raw)
            return (raw/norm)**2
    except TypeError, IOError:
        tol = 1e-9
        def scaling(beta, freq=704.42e6):
            func = list()
            try:
                for bval in beta:
                    if beta_0-tol < bval < beta_0+tol: func.append(1.0)
                    else: func.append((4/(5*pi)) * ((bval**2)/(beta_0**2-bval**2)) * cos(5*pi*beta_0 / (2*bval)))
            except TypeError:
                if beta_0-tol < beta < beta_0+tol: func.append(1.0)
                else: func.append((4/(5*pi)) * ((beta**2)/(beta_0**2-beta**2)) * cos(5*pi*beta_0 / (2*beta)))
            return array(func)**2
    #def scaling(beta, freq): return [1.0]
    return scaling

if __name__=='__main__':
    beta0 = 0.7
    cavL = 5 * (beta0*(c/704.42e6) / 2)
    #cav1 = cavity(I_b=50e-3, V_cav=3.15224e6, beamph=-13.857502, Ql=8e5, beta_0=beta0, beta=0.556752)
    cav1 = cavity(I_b=50e-3, V_cav=3.15224e6*cavL, beamph=-28.724478866, RoverQ=300.0, Ql=8e5, beta_0=beta0, beta=0.556752)
    #cav1 = cavity(I_b=50e-3, V_cav=15.262e6*cavL, beamph=-13.985580262, RoverQ=300.0, Ql=nan, beta_0=beta0, beta=beta0)
    print "Power from the generator = %f kW" % (cav1.Pg/1e3)
    print "Power to the beam = %f kW" % (cav1.Pbeam/1e3)
    print "Implied impedance = %e Ohms" % ((abs(cav1.V_cav)**2) / cav1.Pbeam)
    print "Ql = %e"% cav1.Ql
    print "Scaled (R/Q)Ql = %e Ohms" % (cav1.RoverQ * cav1.RQscaling * cav1.Ql)
    print "R/Q scaling = %f" % cav1.RQscaling
    print "(R/Q).Ql.Ib.cos(phib) / 2.V_cav = %f" % (cav1.RoverQ*cav1.Ql*abs(cav1.I_b)*cos(cav1.beamph*(pi/180)) / (2*abs(cav1.V_cav)))
    print "Scaled (R/Q).Ql.Ib.cos(phib) / 2.V_cav = %f" % (cav1.RoverQ*cav1.RQscaling*cav1.Ql*abs(cav1.I_b)*cos(cav1.beamph*(pi/180)) / (2*abs(cav1.V_cav)))
    cav1.drawvecs(fig=plt.figure(1))

    plt.show()

